#!/usr/bin/env bash

# Slurm: Change a job's QOS
# Homepage: https://github.com/OleHolmNielsen/Slurm_tools/

# List or set nice values:
listqos=0
# Default QOS 
export QOS=normal
# squeue: the %q option prints QOS
# SQUEUE_FORMAT="%q %Q %T %A %P %u %g %S %e %j %b %R"
SQUEUE_FORMAT=QOS,Priority,State,JobID,Partition,UserName,Groupname,StartTime,EndTime,Name,tres-per-node,Reason

# Parse command options
while getopts "lp:q:h" options; do
	case $options in
		l )	listqos=1
			;;
		p )	export partition="-p $OPTARG"
			echo "Select partition $OPTARG"
			listqos=1
			;;
		q )	export QOS=$OPTARG
			echo Set job QOS to $QOS
			;;
		h | * )	echo "Usage:"
			echo "Listing: $0 -l [-q QOS] [-p partition]"
			echo "Setting: $0 [-q QOS] job [jobs ...]"
			echo "Available QOS'es are:"
			sacctmgr --noheader show qos format=name
			exit 1;;
	esac
done
shift $((OPTIND-1))

if test $listqos -eq 0
then
	JOBLIST="$*"
	if test -z "$JOBLIST"
	then
		echo $0: No job list provided
		$0 -h
		exit 1
	elif ! [[ ${JOBLIST} =~ ^[0-9_,]+$ ]]
	then
		echo ERROR: Joblist does not look like jobids:
		echo $JOBLIST
		exit 1
	fi
	echo Job list: $JOBLIST
	# Change QOS of jobs
	for j in $JOBLIST
	do
		scontrol update jobid=$j qos=$QOS
	done
	# Print a job status
	squeue -O "$SQUEUE_FORMAT" -j `echo $JOBLIST | sed 's/ /,/g'`
else
	echo; echo "Jobs with QOS=$QOS:"; echo
	squeue $partition -q $QOS -O "$SQUEUE_FORMAT" 
fi
